package com.tcs.kafka.producer.sender.balance;

import java.util.List;
import java.util.Map;
import java.util.Random;

import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Maps;
import com.tcs.kafka.producer.sender.SenderBalance;

/**
 * 负载均衡器
* @Title: TopicPartition.java
* @Package com.tcs.kafka.producer.sender.balance
* @author 神经刀
* @date 2018年10月15日
* @version V1.0
 */
public class TopicPartition implements ITopicPartition {
	
	private final Logger logger = LoggerFactory.getLogger(TopicPartition.class);
	
	private String zookeeper;
	
	protected CuratorFramework cf;
	
	private Map<String , Integer> cacheMap = Maps.newHashMap();
	
	private Random random = new Random();
	
	private int result = 0;
	
	public TopicPartition() {}
	
	public TopicPartition(String zookeeper) {
		this.zookeeper = zookeeper;
		// 重连策略,没1一秒重试一次,最大重试次数3次
		RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
		cf = CuratorFrameworkFactory.builder().connectString(zookeeper).sessionTimeoutMs(5000).retryPolicy(retryPolicy)
				.build();
		cf.start();
		logger.info(" TopicPartition init : {} ");
	}

	@Override
	public Integer getSizeByTopic(String topic) {
		Integer size = 0;
		try {
			if (!cacheMap.containsKey(topic)) {
				List<String> children = cf.getChildren().forPath("/brokers/topics/" + topic + "/partitions");
				size = children.size();
				cacheMap.put(topic, size);
			} else {
				size = cacheMap.get(topic);
			}
		} catch (Exception e) {
			logger.error("" , e);
		}
		return size;
	}

	@Override
	public Integer getParititon(String topic, SenderBalance SenderBalance) {
		int size = this.getSizeByTopic(topic);
		logger.info(" size : {} ", size);
		switch (SenderBalance) {
		case poll:
			if (result == size - 1) {
				result = 0;
			} else {
				result++;
			}
			break;
		case random:
			result = random.nextInt(size);
			break;
		default:
			break;
		}
		return result;
	}

	public String getZookeeper() {
		return zookeeper;
	}

	public void setZookeeper(String zookeeper) {
		this.zookeeper = zookeeper;
	}
}